from ai.audo_model.s2t.modules import subsampling
import torch
from ai.utils import utils_model


def do_test():
    """"""
    print('耿雪龙')
    model = subsampling.Conv2dSubsampling4(idim=100, odim=200, dropout_rate=0.0)
    input = torch.randn(123, 1000, 100)
    x_mask = torch.ones(123, 1, 1000)
    output, pos_emb, x_mask = model(input, x_mask)
    print(output.shape)
    print(pos_emb.shape)
    print(x_mask.shape)
    num_out = utils_model.get_output_dim_for_conv(utils_model.get_output_dim_for_conv(1000, 3, 2), 3, 2)
    print(f'输入1000维度, 降采样4倍后得到:{num_out}')
    idim = 1000
    print(
        f"{utils_model.get_output_dim_for_conv(utils_model.get_output_dim_for_conv(idim, 3, 2), 5, 3)},,,{(((idim - 1) // 2 - 2) // 3)}")
    conv_output_dim = utils_model.get_output_dim_for_conv(
        utils_model.get_output_dim_for_conv(utils_model.get_output_dim_for_conv(idim, 3, 2), 3, 2), 3, 2)
    conv_output_dim2 = ((((idim - 1) // 2 - 1) // 2 - 1) // 2)
    print(f"{conv_output_dim},{conv_output_dim2}")


def do_test_1():
    num = 2 * (  # 感受野计算
            3 - 1) + 3  # stride_1 * (kernel_size_2 - 1) + kerel_size_1
    num2 = utils_model.get_receptive_field_size_for_continuous_conv([3, 3], [2, 2])
    print(f'{num},{num2}')


def do_test_2():
    """"""
    model = subsampling.Conv2dSubsampling4Pure(idim=100, odim=1, dropout_rate=0.0)
    input = torch.randn(123, 1000, 100)
    lens = torch.randint(0, 100, (123, ))
    output, lens_out = model(input, lens)
    print(output.shape)
    print(lens_out.shape)


if __name__ == "__main__":
    """"""
    do_test_2()
